#include "core/task.h"
#include "tools/klib.h"
#include "tools/log.h"
#include "tools/list.h"
#include "cpu/cpu.h"
#include "cpu/irq.h"
#include "cpu/mmu.h"
#include "os_cfg.h"
#include "core/memory.h"
#include "core/task.h"
#include "core/syscall.h"
#include "comm/elf.h"
#include "comm/cpu_instr.h"
#include "fs/fs.h"


static task_manager_t task_manager;     // 任务管理器
static uint32_t idle_task_stack[IDLE_TASK_STACK_SIZE];  	// 空闲任务堆栈
static task_t task_table[TASK_NR];      // 用户进程表
static mutex_t task_table_mutex;        // 进程表互斥访问锁

static int tss_init(task_t *task, int flag, uint32_t entry, uint32_t esp) {
    uint32_t tss_sel = gdt_alloc_desc();
    if (tss_sel < 0) {
        log_printf("alloc tss failed.");
        return -1;
    }
    
    segment_desc_set(tss_sel, (uint32_t)&task->tss, sizeof(tss_t),
        SEG_P | SEG_DPL0 | SEG_TYPE_TSS);

    kernel_memset(&task->tss, 0, sizeof(tss_t));

    uint32_t kernel_stack = memory_alloc_page();

    if (kernel_stack == 0) {
        goto tss_init_failed;
    }

    int code_sel, data_sel;
    if (flag & TASK_FLAGS_SYSTEM) {
        code_sel = KERNEL_SELECTOR_CS;
        data_sel = KERNEL_SELECTOR_DS;
    } else {
        code_sel = task_manager.app_code_sel | SEG_CPL3;
        data_sel = task_manager.app_data_sel | SEG_RPL3;
    }

    
    task->time_ticks = TASK_TIME_SLICE_DEFAULT;
    task->slice_ticks = task->time_ticks;
    
    task->tss.eip = entry;
    task->tss.esp = esp;
    task->tss.esp0 = kernel_stack + MEM_PAGE_SIZE;
    task->tss.ss = data_sel;
    task->tss.ss0 = KERNEL_SELECTOR_DS;
    task->tss.es = task->tss.ds = task->tss.fs = task->tss.gs = data_sel;
    task->tss.cs = code_sel;
    task->tss.eflags = EFLAGS_IF | EFLAGS_DEFAULT;
    task->tss_sel = tss_sel;

    // 为当前进程分配页目录表
    uint32_t page_dir = memory_create_uvm();
    if (page_dir == 0) {
        goto tss_init_failed;
    }
    // 设置页目录表的物理地址
    task->tss.cr3 = page_dir;

    return 0;
tss_init_failed:
    gdt_free_sel(tss_sel);
    if (kernel_stack) {
        memory_free_page(kernel_stack);
    }
    return -1;
}

int task_init(task_t *task, const char *name, int flag, uint32_t entry, uint32_t esp) {
    ASSERT(task != (task_t *)0);
    
    tss_init(task, flag, entry, esp);

    task->pid = (uint32_t)task;
    task->status = 0;
    task->parent = (task_t *)0;
    kernel_strncpy(task->name, name, TASK_NAME_SIZE);

    task->heap_start = 0;
    task->heap_end = 0;

    task->state = TASK_CREATED;
    task->sleep_ticks = 0;

    list_node_init(&task->all_node);
    list_node_init(&task->run_node);
    list_node_init(&task->wait_node);

    kernel_memset(&task->file_table, 0, sizeof(file_t *) * TASK_OFILE_NR);

    irq_state_t state = irq_enter_protection();
    list_insert_last(&task_manager.task_list, &task->all_node);
    irq_leave_protection(state);

    return 0;
}

void task_start(task_t *task) {
    irq_state_t state = irq_enter_protection();
    task_set_ready(task);
    irq_leave_protection(state);
}

// 回退task_init操作的产生的影响
void task_uninit(task_t *task) {
    if (task->tss_sel) {
        gdt_free_sel(task->tss_sel);
    }

    if (task->tss.esp0) {
        memory_free_page(task->tss.esp - MEM_PAGE_SIZE);
    }

    if (task->tss.cr3) {
        memory_destory_uvm(task->tss.cr3);
    }

    kernel_memset(task, 0, sizeof(task));
}


// void simple_switch(uint32_t **from, uint32_t *to);
    

void task_switch_from_to(task_t *from, task_t *to) {
    switch_to_tss(to->tss_sel);
    // simple_switch(&from->stack, to->stack);
}

void task_time_tick(void) {
    task_t *curr_task = task_current();
    // 时间片的处理
    irq_state_t state = irq_enter_protection();
    if (--curr_task->slice_ticks == 0) {
        curr_task->slice_ticks = curr_task->time_ticks;

        task_set_block(curr_task);
        task_set_ready(curr_task);
    }

    list_node_t *curr = list_first(&task_manager.sleep_list);
    while (curr) {
        list_node_t *next = list_node_next(curr);
        task_t *task = list_node_parent(curr, task_t, run_node);
        if (--task->sleep_ticks == 0) {
            task_set_wakeup(task);
            task_set_ready(task);
        }
        curr = next;
    }
    task_dispatch();
    irq_leave_protection(state);
}

void task_first_init(void) {
    void first_task_entry(void);
    extern uint8_t s_first_task[], e_first_task[];

    // 第一个进程在内存中占用的大小
    uint32_t copy_size = (uint32_t)(e_first_task - s_first_task);
    // 分配一个比copy_size更大一些的空间，多出来的空间可以用作进程的栈和堆
    uint32_t alloc_size = 10 * MEM_PAGE_SIZE;
    ASSERT(copy_size <= alloc_size);

    uint32_t first_start = (uint32_t)first_task_entry;

    task_init(&task_manager.first_task, "first task", TASK_FLAGS_USER, first_start, first_start + alloc_size);
    task_manager.first_task.heap_start = (uint32_t)e_first_task;
    task_manager.first_task.heap_end = (uint32_t)e_first_task;
    write_tr(task_manager.first_task.tss_sel);
    task_manager.curr_task = &task_manager.first_task;

    // 切换到当前进程的页目录表(之前使用的是操作系统的页目录表kernel_page_dir)
    mmu_set_page_dir(task_manager.first_task.tss.cr3);

    // 给进程分配alloc_size大小的虚拟空间
    memory_alloc_page_for(first_start, alloc_size, PTE_P | PTE_W | PTE_U);
    kernel_memcpy((void *)first_start, s_first_task, copy_size);
    task_start(&task_manager.first_task);
}

task_t * task_first_task(void) {
    return &task_manager.first_task;
}

static void idle_task_entry(void) {
    while (1) { hlt(); }
}

void task_manager_init(void) {
    kernel_memset(task_table, 0, sizeof(task_table));
    mutex_init(&task_table_mutex);

    int sel = gdt_alloc_desc();
    segment_desc_set(sel, 0x00000000, 0xFFFFFFFF, 
    SEG_P  | SEG_DPL3 | SEG_S_NORMAL | SEG_TYPE_DATA | SEG_TYPE_DATA_RW | SEG_D);
    task_manager.app_data_sel = sel;

    sel = gdt_alloc_desc();
    segment_desc_set(sel, 0x00000000, 0xFFFFFFFF, 
    SEG_P  | SEG_DPL3 | SEG_S_NORMAL | SEG_TYPE_CODE | SEG_TYPE_CODE_RA | SEG_D);
    task_manager.app_code_sel = sel;


    list_init(&task_manager.ready_list);
    list_init(&task_manager.task_list);
    list_init(&task_manager.sleep_list);
    task_manager.curr_task = (task_t *)0;

    task_init(&task_manager.idle_task, "idle task", TASK_FLAGS_SYSTEM, 
        (uint32_t)idle_task_entry, 
        (uint32_t)(idle_task_stack + IDLE_TASK_STACK_SIZE));
    task_start(&task_manager.idle_task);
}

void task_set_ready(task_t *task) {
    if (task == &task_manager.idle_task) {
        return;
    }

    list_insert_last(&task_manager.ready_list, &task->run_node);
    task->state = TASK_READY;
}

void task_set_block(task_t *task) {
    if (task == &task_manager.idle_task) {
        return;
    }

    list_remove(&task_manager.ready_list, &task->run_node);
}

task_t * task_current() {
    return task_manager.curr_task;
}

task_t * task_next_run() {
    if (list_count(&task_manager.ready_list) == 0) {
        return &task_manager.idle_task;
    }
    list_node_t *task_node = list_first(&task_manager.ready_list);
    return list_node_parent(task_node, task_t, run_node);
}

int sys_yield(void) {
    irq_state_t state = irq_enter_protection();
    if (list_count(&task_manager.ready_list) > 1) {
        task_t * curr_task = task_current();
        
        task_set_block(curr_task);
        task_set_ready(curr_task);

        task_dispatch();
    }
    irq_leave_protection(state);

    return 0;
}

void task_dispatch(void) {
    irq_state_t state = irq_enter_protection();
    task_t *to = task_next_run();
    if (to != task_manager.curr_task) {
        task_t *from = task_current();
        task_manager.curr_task = to;
        // to->state = TASK_RUNNING;
        // from->state = TASK_READY;
        task_switch_from_to(from, to);
    }
    irq_leave_protection(state);
}

void task_set_sleep(task_t *task, uint32_t ticks) {
    if (ticks == 0) {
        return;
    }

    task->sleep_ticks = ticks;
    task->state = TASK_SLEEP;
    list_insert_last(&task_manager.sleep_list, &task->run_node);
}

void task_set_wakeup(task_t *task) {
    list_remove(&task_manager.sleep_list, &task->run_node);
}

file_t *task_file(int fd) {
    if (fd >= 0 && fd < TASK_OFILE_NR) {
        file_t *file = task_current()->file_table[fd];
        return file;
    }
    return (file_t *)0;
}

int task_alloc_fd(file_t *file) {
    task_t *task = task_current();
    for (int i = 0; i < TASK_OFILE_NR; i++) {
        file_t *p = task->file_table[i];
        if (p == (file_t *)0) {
            task->file_table[i] = file;
            return i;
        }      
    }
    return -1;
}

void task_remove_fd(int fd) {
    if (fd >= 0 && fd < TASK_OFILE_NR) {
        task_current()->file_table[fd] = (file_t *)0;
    }
}

void sys_sleep(uint32_t ms) {
    irq_state_t state = irq_enter_protection();

    task_set_block(task_manager.curr_task);
    task_set_sleep(task_manager.curr_task, (ms + (OS_TICKS_MS - 1)) / OS_TICKS_MS);

    task_dispatch();

    irq_leave_protection(state);
}

int sys_getpid(void) {
    return task_current()->pid;
}

static task_t * alloc_task(void) {
    task_t *task = (task_t *)0;
 
    mutex_lock(&task_table_mutex);
    for (int i = 0; i < TASK_NR; i++) {
        task_t *curr = task_table + i;
        if (curr->name[0] == '\0') {
            task = curr;
            break;
        }
    }
    mutex_unlock(&task_table_mutex);
 
    return task;
}

static void free_task(task_t *task) {
    mutex_lock(&task_table_mutex);
    task->name[0] = '\0';
    mutex_unlock(&task_table_mutex);
}

static void copy_opened_files(task_t *child_task) {
    task_t *parent = task_current();

    for (int i = 0; i < TASK_OFILE_NR; i++) {
        file_t *file = parent->file_table[i];
        if (file) {
            file_inc_ref(file);
            child_task->file_table[i] = file;
        }
    }
}

int sys_fork(void) {
    task_t *parent_task = task_current();
    task_t *child_task = alloc_task();
    if (child_task == (task_t *)0) {
        goto fork_failed;
    }

    syscall_frame_t *frame = (syscall_frame_t *)(parent_task->tss.esp0 - sizeof(syscall_frame_t));
    int err = task_init(child_task, parent_task->name, TASK_FLAGS_USER, frame->eip, frame->esp + sizeof(uint32_t) * SYSCALL_PARM_COUNT);
    if (err < 0) {
        goto fork_failed;
    }

    copy_opened_files(child_task);

    tss_t *tss = &child_task->tss;
    tss->eax = 0;               //子进程返回值
    tss->ebx = frame->ebx;
    tss->ecx = frame->ecx;
    tss->edx = frame->edx;
    tss->esi = frame->esi;
    tss->edi = frame->edi;
    tss->ebp = frame->ebp;

    tss->cs = frame->cs;
    tss->ds = frame->ds;
    tss->es = frame->es;
    tss->fs = frame->fs;
    tss->gs = frame->gs;
    
    tss->eflags = frame->eflags;
    child_task->parent = parent_task;
    
    // 为子进程分配空间并把父进程的各个代码段都内容复制到子进程新分配到内存空间中
    // uvm user virtual memory
    if ((tss->cr3 = memory_copy_uvm(parent_task->tss.cr3)) == 0) {
        goto fork_failed;
    }

    task_start(child_task);

    return child_task->pid;     // 父进程返回值
fork_failed:

    if (child_task) {
        task_uninit(child_task);
        free_task(child_task);
    }
    return -1;
}

// 获取指定虚拟地址的物理地址
uint32_t memory_get_paddr (uint32_t page_dir, uint32_t vaddr) {
    pte_t * pte = find_pte((pde_t *)page_dir, vaddr, 0);
    if (pte == (pte_t *)0) {
        return 0;
    }

     return pte_addr(pte) + (vaddr & (MEM_PAGE_SIZE - 1));
}


// 加载一个程序表头的数据到内存中
static int load_phdr(int file, Elf32_Phdr * phdr, uint32_t page_dir) {
    // 生成的ELF文件要求是页边界对齐的(gcc默认会对vaddr进行4k对齐)
    ASSERT((phdr->p_vaddr & (MEM_PAGE_SIZE - 1)) == 0);

    // 为每个需要加载到内存中的段（.txt .rodata .data .bss）分配空间（包括bss段，主要但是我们没有初始化这个bss段）
    int err = memory_alloc_for_page_dir(page_dir, phdr->p_vaddr, phdr->p_memsz, PTE_P | PTE_U | PTE_W);
    if (err < 0) {
        log_printf("no memory");
        return -1;
    }

    // 调整当前的读写位置
    if (sys_lseek(file, phdr->p_offset, 0) < 0) {
        log_printf("read file failed");
        return -1;
    }

    // 为段分配所有的内存空间.后续操作如果失败，将在上层释放
    // 简单起见，设置成可写模式，也许可考虑根据phdr->flags设置成只读
    // 因为没有找到该值的详细定义，所以没有加上
    uint32_t vaddr = phdr->p_vaddr;
    uint32_t size = phdr->p_filesz;
    while (size > 0) {
        int curr_size = (size > MEM_PAGE_SIZE) ? MEM_PAGE_SIZE : size;

        uint32_t paddr = memory_get_paddr(page_dir, vaddr);

        // 注意，这里用的页表仍然是当前的
        if (sys_read(file, (char *)paddr, curr_size) <  curr_size) {
            log_printf("read file failed");
            return -1;
        }

        size -= curr_size;
        vaddr += curr_size;
    }

    // bss区考虑由crt0和cstart自行清0，这样更简单一些
    // 如果在上边进行处理，需要考虑到有可能的跨页表填充数据，懒得写代码
    // 或者也可修改memory_alloc_for_page_dir，增加分配时清0页表，但这样开销较大
    // 所以，直接放在cstart的crt0中直接内存填0，比较简单
    return 0;
}

static uint32_t load_elf_file(task_t *task, const char *name, uint32_t page_dir) {
    Elf32_Ehdr elf_hdr;
    Elf32_Phdr elf_phdr;
    int file = sys_open(name, 0);
    if (file < 0) {
        log_printf("open failed. %s", name);
        goto load_failed;
    }    

    int cnt = sys_read(file, (char *)&elf_hdr, sizeof(elf_hdr));
    if (cnt < sizeof(Elf32_Ehdr)) {
        log_printf("elf hdr too samll. size=%d", cnt);
        goto load_failed;
    }

    // 做点必要性的检查。当然可以再做其它检查
    if ((elf_hdr.e_ident[0] != ELF_MAGIC) || (elf_hdr.e_ident[1] != 'E')
        || (elf_hdr.e_ident[2] != 'L') || (elf_hdr.e_ident[3] != 'F')) {
        log_printf("check elf indent failed.");
        goto load_failed;
    }

    // 必须是可执行文件和针对386处理器的类型，且有入口
    if ((elf_hdr.e_type != ET_EXEC) || (elf_hdr.e_machine != ET_386) || (elf_hdr.e_entry == 0)) {
        log_printf("check elf type or entry failed.");
        goto load_failed;
    }

    // 必须有程序头部
    if ((elf_hdr.e_phentsize == 0) || (elf_hdr.e_phoff == 0)) {
        log_printf("none programe header");
        goto load_failed;
    }

    uint32_t e_phoff = elf_hdr.e_phoff;
    for (int i = 0; i < elf_hdr.e_phnum; i++, e_phoff += elf_hdr.e_phentsize) {
        if (sys_lseek(file, e_phoff, 0) < 0) {
            log_printf("read file failed.");
            goto load_failed;
        }

        cnt = sys_read(file, (char *)&elf_phdr, sizeof(elf_phdr));
        if (cnt < sizeof(elf_phdr)) {
            log_printf("read file failed.");
            goto load_failed;
        }
    
        // 主要判断是否时可加载的类型段，并且要求加载的地址必须是用户空间
        if ((elf_phdr.p_type != PT_LOAD) || (elf_phdr.p_vaddr < MEMORY_TASK_BASE)) {
            continue;
        }
        // 加载当前程序头所指明的段
        int err = load_phdr(file, &elf_phdr, page_dir);
        if (err < 0) {
            log_printf("load program hdr failed");
            goto load_failed;
        }
        task->heap_start = task->heap_end = elf_phdr.p_vaddr + elf_phdr.p_memsz;
    }

    sys_close(file);
    return elf_hdr.e_entry;

load_failed:
    if (file) {
        sys_close(file);
    }
    return 0;
}

static int copy_args(char *to, uint32_t page_dir, int argc, char **argv) {
    task_args_t task_args;
    task_args.argc = argc;
    task_args.argv = (char **)(to + sizeof(task_args_t));

    char *dest_arg = to + sizeof(task_args_t) + sizeof(char *) * (argc + 1);
    char **dest_argv_tb = (char **)memory_get_paddr(page_dir, (uint32_t)(to + sizeof(task_args_t)));
    ASSERT(dest_argv_tb != 0);

    for (int i = 0; i < argc; i++) {
        char *from = argv[i];
        int len = kernel_strlen(from) + 1;
        int err = memory_copy_uvm_data((uint32_t)dest_arg, page_dir, (uint32_t)from, len);
        ASSERT(err >= 0);

        dest_argv_tb[i] = dest_arg;
        dest_arg += len;
    }
        // 可能存在无参的情况，此时不需要写入
    if (argc) {
        dest_argv_tb[argc] = '\0';
    }
    return memory_copy_uvm_data((uint32_t)to, page_dir, (uint32_t)&task_args, sizeof(task_args));
}

int sys_execve(char *name, char **argv, char **env) {
    task_t *task = task_current();

    kernel_strncpy(task->name, get_file_name(name), TASK_NAME_SIZE);

    uint32_t old_page_dir = task->tss.cr3;
    uint32_t new_page_dir = memory_create_uvm();
    if (! new_page_dir) {
        goto exec_failed;
    }
    
    uint32_t entry = load_elf_file(task, name, new_page_dir);
    if (entry == 0) {
        goto exec_failed;
    }
    
    uint32_t stack_top = MEM_TASK_STACK_TOP - MEM_TASK_ARG_SIZE;
    int err = memory_alloc_for_page_dir(new_page_dir, MEM_TASK_STACK_TOP - MEM_TASK_STACK_SIZE,
        MEM_TASK_STACK_SIZE, PTE_P | PTE_U | PTE_W);

    if (err < 0) {
        goto exec_failed;
    }

    int argc = strings_count(argv);
    err = copy_args((char *)stack_top, new_page_dir, argc, argv);
    if (err < 0) {
        goto exec_failed;
    }
    // 加载完毕，为程序的执行做必要准备
    // 注意，exec的作用是替换掉当前进程，所以只要改变当前进程的执行流即可
    // 当该进程恢复运行时，像完全重新运行一样，所以用户栈要设置成初始模式
    // 运行地址要设备成整个程序的入口地址
    syscall_frame_t *frame = (syscall_frame_t *)(task->tss.esp0 - sizeof(syscall_frame_t));
    frame->eip = entry;
    frame->eax = frame->ebx = frame->ecx = frame->edx = 0;
    frame->esi = frame->edi = frame->ebp = 0;
    frame->eflags = EFLAGS_DEFAULT| EFLAGS_IF;  // 段寄存器无需修改

    // 内核栈不用设置，保持不变，后面调用memory_destroy_uvm并不会销毁内核栈的映射。
    // 但用户栈需要更改, 同样要加上调用门的参数压栈空间
    frame->esp = stack_top - sizeof(uint32_t)*SYSCALL_PARM_COUNT;

    task->tss.cr3 = new_page_dir;
    mmu_set_page_dir(new_page_dir);
    memory_destory_uvm(old_page_dir);

    return 0;

exec_failed:
    if (new_page_dir) {
        task->tss.cr3 = old_page_dir;
        mmu_set_page_dir(old_page_dir);
        memory_destory_uvm(new_page_dir);
    }
    return -1;
}

void sys_exit(int status) {
    task_t *curr_task = task_current();
    for (int fd = 0; fd < TASK_OFILE_NR; fd++) {
        file_t *file = curr_task->file_table[fd];
        if (file) {
            sys_close(fd);
            curr_task->file_table[fd] = (file_t *)0;
        }
    }

    int moved_child_zoombie = 0;
    mutex_lock(&task_table_mutex);
    for (int i = 0; i < TASK_NR; i++) {
        task_t *task = task_table + i;
        if (task->parent == curr_task) {
            task->parent = &task_manager.first_task;
            if (task->state == TASK_ZOOMBIE) {
                moved_child_zoombie = 1;
            }
        }
    } 
    mutex_unlock(&task_table_mutex);

    irq_state_t state = irq_enter_protection();

    task_t *parent = curr_task->parent;

    // 如果有移动子进程，则唤醒init进程
    if (moved_child_zoombie && (parent != &task_manager.first_task)) {
        if (task_manager.first_task.state == TASK_WAITTING) {
            task_set_ready(&task_manager.first_task);
        }
    }

    // 如果有父任务在wait，则唤醒父任务进行回收
    // 如果父进程没有等待，则一直处理僵死状态
    if (parent->state == TASK_WAITTING) {
        task_set_ready(curr_task->parent);
    }

    curr_task->state = TASK_ZOOMBIE;
    curr_task->status = status;
    task_set_block(curr_task);
    task_dispatch();
    irq_leave_protection(state);
}

int sys_wait(int *status) {
    task_t *curr_task = task_current();
    while (1) {
        mutex_lock(&task_table_mutex);
        for (int i = 0; i < TASK_NR; i++) {
            task_t *task = task_table + i;
            if (task->parent != curr_task) {
                continue;
            }

            if (task->state == TASK_ZOOMBIE) {
                int pid = task->pid;
                *status = task->status;

                memory_destory_uvm(task->tss.cr3);
                memory_free_page(task->tss.esp0 - MEM_PAGE_SIZE);
                kernel_memset(task, 0, sizeof(task_t));

                mutex_unlock(&task_table_mutex);
                return pid;
            }
        } 

        mutex_unlock(&task_table_mutex);

        // 找不到子进程说明子进程还在执行，则等待。等子进程调用_exit函数时把父进程唤醒   
        irq_state_t state = irq_enter_protection();
        task_set_block(curr_task);
        curr_task->state = TASK_WAITTING;
        task_dispatch();
        irq_leave_protection(state);
    }
}